//! Trie data structure implementation in Rust.
//! This module provides a generic Trie that can store key-value pairs,
//! with operations for insertion, lookup, removal, and replacement.

use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::Hash;

/// A generic Trie data structure.
///
/// # Type Parameters
/// - `K`: The key type, which must implement `Eq`, `Hash`, and `Clone`.
/// - `V`: The value type, which must implement `Clone`.
///
/// # Examples
/// ```
/// use trie::Trie;
///
/// let mut trie = Trie::new();
/// trie.insert("hello", "world");
/// assert_eq!(trie.lookup("hello"), Some(&"world"));
/// ```
pub struct Trie<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    root: TrieNode<K, V>,
    entries: usize,
}

/// A node in the Trie, containing child nodes and an optional value.
struct TrieNode<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    children: HashMap<K, TrieNode<K, V>>,
    value: Option<V>,
}

impl<K, V> Trie<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    /// Creates a new, empty Trie.
    pub fn new() -> Self {
        Trie {
            root: TrieNode::new(),
            entries: 0,
        }
    }

    /// Inserts a key-value pair into the Trie.
    ///
    /// # Arguments
    /// - `key`: The key to insert.
    /// - `value`: The value to associate with the key.
    ///
    /// # Returns
    /// - `true` if the insertion was successful, `false` otherwise.
    pub fn insert(&mut self, key: &[K], value: V) -> bool {
        let mut node = &mut self.root;
        for k in key {
            node = node.children.entry(k.clone()).or_insert_with(TrieNode::new);
        }
        if node.value.is_none() {
            self.entries += 1;
        }
        node.value = Some(value);
        true
    }

    /// Looks up a value by its key.
    ///
    /// # Arguments
    /// - `key`: The key to search for.
    ///
    /// # Returns
    /// - `Some(&V)` if the key exists, `None` otherwise.
    pub fn lookup(&self, key: &[K]) -> Option<&V> {
        let mut node = &self.root;
        for k in key {
            if let Some(child) = node.children.get(k) {
                node = child;
            } else {
                return None;
            }
        }
        node.value.as_ref()
    }

    /// Removes a key-value pair from the Trie.
    ///
    /// # Arguments
    /// - `key`: The key to remove.
    ///
    /// # Returns
    /// - `true` if the key was removed, `false` if it did not exist.
    pub fn remove(&mut self, key: &[K]) -> bool {
        let mut node = &mut self.root;
        for k in key {
            if let Some(child) = node.children.get_mut(k) {
                node = child;
            } else {
                return false;
            }
        }
        if node.value.is_some() {
            node.value = None;
            self.entries -= 1;
            true
        } else {
            false
        }
    }

    /// Returns the number of entries in the Trie.
    pub fn num_entries(&self) -> usize {
        self.entries
    }
}

impl<K, V> TrieNode<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    /// Creates a new, empty TrieNode.
    fn new() -> Self {
        TrieNode {
            children: HashMap::new(),
            value: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Trie;

    const NUM_TEST_VALUES: usize = 10000;

    /// Generates a Trie filled with a large number of values.
    fn generate_trie() -> Trie<char, i32> {
        let mut trie = Trie::new();
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            trie.insert(&key.chars().collect::<Vec<_>>(), i as i32);
        }
        trie
    }

    #[test]
    fn test_trie_new_free() {
        let mut trie = Trie::new();
        assert_eq!(trie.num_entries(), 0);

        trie.insert(&['h', 'e', 'l', 'l', 'o'], "there");
        trie.insert(&['h', 'e', 'l', 'l'], "testing");
        trie.insert(&['t', 'e', 's', 't', 'i', 'n', 'g'], "testing");
        trie.insert(&[], "asfasf");
        assert_eq!(trie.num_entries(), 4);

        trie.remove(&['h', 'e', 'l', 'l', 'o']);
        assert_eq!(trie.num_entries(), 3);
    }

    #[test]
    fn test_trie_insert() {
        let mut trie = generate_trie();
        let entries = trie.num_entries();

        // Inserting a NULL value should have no effect.
        assert!(!trie.insert(&['h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd'], 0));
        assert_eq!(trie.num_entries(), entries);
    }

    #[test]
    fn test_trie_lookup() {
        let trie = generate_trie();

        // Lookup for non-existent values.
        assert_eq!(trie.lookup(&['0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0']), None);
        assert_eq!(trie.lookup(&[]), None);

        // Look up all values.
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            let key_chars: Vec<_> = key.chars().collect();
            assert_eq!(trie.lookup(&key_chars), Some(&(i as i32)));
        }
    }

    #[test]
    fn test_trie_remove() {
        let mut trie = generate_trie();
        let entries = trie.num_entries();

        // Remove non-existent values.
        assert!(!trie.remove(&['0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0']));
        assert!(!trie.remove(&[]));

        // Remove all values.
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            let key_chars: Vec<_> = key.chars().collect();
            assert!(trie.remove(&key_chars));
            assert_eq!(trie.num_entries(), entries - i - 1);
        }
    }

    #[test]
    fn test_trie_replace() {
        let mut trie = generate_trie();
        let val = 999;
        let key: String = format!("{}", val);
        let key_chars: Vec<_> = key.chars().collect();

        trie.insert(&key_chars, val);
        assert_eq!(trie.num_entries(), NUM_TEST_VALUES);
        assert_eq!(trie.lookup(&key_chars), Some(&val));
    }

    #[test]
    fn test_trie_insert_empty() {
        let mut trie = Trie::new();
        let buf = "test";
        trie.insert(&[], buf);
        assert_ne!(trie.num_entries(), 0);
        assert_eq!(trie.lookup(&[]), Some(&buf));
        assert!(trie.remove(&[]));
        assert_eq!(trie.num_entries(), 0);
    }
}
